﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.IdentityModel.Tokens.Saml11;
using Microsoft.IdentityModel.Tokens.Saml2;
using Microsoft.IdentityModel.Tokens;
using System.IdentityModel.Tokens;
using System.Security.Cryptography;
using System.Xml;
using System.IO;
using System.IdentityModel.Selectors;
using Microsoft.IdentityModel.Web;
using Microsoft.IdentityModel.Configuration;
using System.Security.Cryptography.X509Certificates;
using System.Web;
using System.Configuration;

namespace AccessControlLibrary
{
    /// <summary>
    /// EncryptedSecurityTokenHandler with support for certificates stored in a file or a blob in Azure Storage
    /// </summary>
    /// <remarks><p>The location of the certificate must be specified in a custom configuration element :<br/ >
    /// &lt;certificate locationType="..." ... /&gt;. The <b>locationType</b> can be "File", "AzureBlob" or "Value".<br />
    /// "Value" is used to read the certificate from the configuration file. You need to have two more attributes : <b>value</b>, which contains the base64 encoded
    /// certificate, and <b>password</b>, which contains the password used to encrypt the private key.<br />
    /// TODO : "Blob" is used to retrieve the certificate from a blob in Azure Storage.<br />
    /// See sample code.</p>
    /// <p>When you create a scope in Access Control, you can set a certificate to encrypt the security tokens returned by ACS.
    /// ACS uses the public key of the certificate, so your application needs to have access to the private key to decrypt the tokens.
    /// This is normally achieved by specifying a &lt;certificatereference&gt; element in the configuration file, which Geneva Framework
    /// will resolve to find the correct certificate in one of the server's certficate stores.<br />Unfortunately, you cannot access the
    /// certificate stores in Windows Azure, so we need an alternate way to get the certificates.</p>
    /// </remarks>
    class CertificateEncryptedSecurityTokenHandler : EncryptedSecurityTokenHandler
    {
        // Will be used to store the SecurityTokenResolver configured with the certificate
        SecurityTokenResolver _serviceTokenResolver;

        public CertificateEncryptedSecurityTokenHandler()
        {
            _serviceTokenResolver = GetServiceTokenResolver();
        }

        // While it would be tempting and logical to call the base contructor passing the result of GetServiceTokenResolver as a parameter, this doesn't work :
        // Unfortunately, when Geneva Framework reads the configuration from web.config, it will set ServiceTokenResolver to a new value. We therefore store
        // our ServiceTokenResolver in a member, and set it back again in ReadToken method
        public CertificateEncryptedSecurityTokenHandler(XmlElement customConfigElement)
            : base()
        {
            // Creates the default SecurityTokenResolver
            _serviceTokenResolver = GetServiceTokenResolver(customConfigElement);
        }

        /// <summary>
        /// Generates a SecurityTokenResolver with a certificate
        /// </summary>
        /// <param name="customConfigElement">A &lt;certificate&gt; config element used to retrieve the certificate</param>
        /// <returns></returns>
        private static SecurityTokenResolver GetServiceTokenResolver()
        {
            if (ConfigurationManager.AppSettings["certificate"] == null)
            {
                throw new InvalidOperationException("The CertificateEncryptedSecurityTokenHandler class expects the following application setting <add key=\"certificate\" value=\"...\" />.");
            }
            if (ConfigurationManager.AppSettings["certificatePassword"] == null)
            {
                throw new InvalidOperationException("The TripleDESSessionSecurityTokenHandler class expects the following application setting <add key=\"certificatePassword\" value=\"...\" />.");
            }

            X509Certificate2 cert;
            byte[] certBytes;
            try
            {
                certBytes = Convert.FromBase64String(ConfigurationManager.AppSettings["certificate"]);
            }
            catch (Exception ex)
            {
                throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : The \"certificate\" application setting does not contain a valid base64 string.", ex);
            }
            try
            {
                // Testing : adding PersistKeySet flag to avoid occasional errors "The private key is not present in the X.509 certificate"
                // I keep MachineKeySet because I have doubts about storing info in the user's profile.
                cert = new X509Certificate2(certBytes, ConfigurationManager.AppSettings["certificatePassword"], X509KeyStorageFlags.MachineKeySet | X509KeyStorageFlags.PersistKeySet);
            }
            catch (Exception ex)
            {
                throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : Error reading the certificate from the \"certificate\" application setting.", ex);
            }

            List<SecurityToken> list = new List<SecurityToken>();
            list.Add(new X509SecurityToken(cert));
            return SecurityTokenResolver.CreateDefaultSecurityTokenResolver(list.AsReadOnly(), false);
        }

        /// <summary>
        /// Generates a SecurityTokenResolver with a certificate
        /// </summary>
        /// <param name="customConfigElement">A &lt;certificate&gt; config element used to retrieve the certificate</param>
        /// <returns></returns>
        private static SecurityTokenResolver GetServiceTokenResolver(XmlElement customConfigElement)
        {
            if (customConfigElement.LocalName != "certificate")
            {
                throw new InvalidOperationException("The CertificateEncryptedSecurityTokenHandler class expects a custom config element called <certificate>.");
            }
            if (customConfigElement.Attributes["location"] == null)
            {
                throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : You must create a \"location\" attribute in the <certificate> element. Possible values : \"Value\" or \"Blob\"");
            }
            if (customConfigElement.Attributes["password"] == null)
            {
                throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : You must specify the password used to encrypt the private key in the \"password\" attribute.");
            }

            X509Certificate2 cert;
            switch (customConfigElement.Attributes["location"].Value)
            {
                case "Value" :
                    // The certificate is passed as a base64 encoded string in the value attribute
                    if (customConfigElement.Attributes["value"] == null)
                    {
                        throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : You must store the base64-encoded certificate in the \"value\" attribute.");
                    }
                    byte[] certBytes;
                    try
                    {
                        certBytes = Convert.FromBase64String(customConfigElement.Attributes["value"].Value);
                    }
                    catch (Exception ex)
                    {
                        throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : The \"value\" attribute does not contain a valid base64 string.", ex);
                    }
                    try
                    {
                        // Testing : adding PersistKeySet flag to avoid occasional errors "The private key is not present in the X.509 certificate"
                        // I keep MachineKeySet because I have doubts about storing info in the user's profile.
                        cert = new X509Certificate2(certBytes, customConfigElement.Attributes["password"].Value, X509KeyStorageFlags.MachineKeySet | X509KeyStorageFlags.PersistKeySet);
                    }
                    catch (Exception ex)
                    {
                        throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : Error reading the certificate from the \"value\" attibute.", ex);
                    }
                    break;
                case "Blob" :
                    // Could be used to retrieve the certificate from a blob, not implemented yet.
                    // Possible security issue transmitting the certificate from storage to webrole ?
                    throw new NotImplementedException("CertificateEncryptedSecurityTokenHandler : Storing certificates in Azure blobs is not implemented yet.");
                default:
                    throw new InvalidOperationException("CertificateEncryptedSecurityTokenHandler : the locationType attribute value must be \"Value\", \"AzureBlob\" or \"File\"");
            }

            List<SecurityToken> list = new List<SecurityToken>();
            list.Add(new X509SecurityToken(cert));
            return SecurityTokenResolver.CreateDefaultSecurityTokenResolver(list.AsReadOnly(), false);
        }

        public override SecurityToken ReadToken(XmlReader reader)
        {
            // ServiceTokenResolver is changed after the constructor is called, when Geneva Framework applies the configuration, so we need to set it again here
            Configuration.ServiceTokenResolver = _serviceTokenResolver;
            return base.ReadToken(reader);
        }
    }
}
